# **************************************
# --*-- coding: utf-8 --*--
# @Author  : white
# @FileName: tools.py
# @Time    : 2025-08-14
# **************************************
from operator import itemgetter
from pathlib import Path
from typing import List, Tuple, Callable

import matplotlib.pyplot as plt

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False  # 解决坐标轴负号显示问题


def L2_dist(x: Tuple, y: Tuple):
    """
    L2距离函数
    :param x:
    :param y:
    :return:
    """
    return sum(map(lambda u, v: (u - v) ** 2, x, y))


def normalized_dist(x: Tuple, y: Tuple):
    """
    归一化距离函数
    :param x:
    :param y:
    :return:
    """
    # x值保持不变，y值放大10倍
    # 原因是因为在穷举生成空间上的点时，i值缩小10倍，j值缩小100倍，为了保持一致
    return (x[0] - y[0]) ** 2 + (10 * x[1] - 10 * y[1]) ** 2


def draw_scatter(img: Path, title: str, class_result: List = None, **kwargs):
    """
    画散点图
    :param class_result: 分类结果
    :param img: 保存路径
    :param title: 画图标题
    :param kwargs: 格外参数
    :return:
    """
    if class_result is not None:
        # 把每个颜色的数据分开
        yellow = []
        green = []
        blue = []
        for pair in class_result:
            if pair[1] == 'y':
                yellow.append(pair[0])
            if pair[1] == 'g':
                green.append(pair[0])
            if pair[1] == 'b':
                blue.append(pair[0])
        # 将空间中穷举的每个点进行画图，按照不同的颜色进行区分
        plt.scatter([x[0] for x in yellow],
                    [x[1] for x in yellow],
                    s=1, c='yellow', marker='1', alpha=0.3)
        plt.scatter([x[0] for x in blue],
                    [x[1] for x in blue],
                    s=1, c='blue', marker='2', alpha=0.3)
        plt.scatter([x[0] for x in green],
                    [x[1] for x in green],
                    s=1, c='green', marker='3', alpha=0.3)

    if kwargs.get("point_x", None) is not None and kwargs.get("point_y", None) is not None:
        plt.scatter(kwargs["point_x"], kwargs["point_y"], c='r', s=200, marker='*', alpha=0.8, zorder=10)

    plt.scatter(kwargs["grief_heights"], kwargs["grief_weights"], c='g', marker='s', s=50, alpha=0.8, zorder=10)
    plt.scatter(kwargs["agony_heights"], kwargs["agony_weights"], c='b', marker='^', s=50, alpha=0.8, zorder=10)
    plt.scatter(kwargs["despair_heights"], kwargs["despair_weights"], c='y', s=50, alpha=0.8, zorder=10)

    # 设置坐标轴范围 (xmin, xmax, ymin, ymax)
    plt.axis((10, 70, 1, 7))
    plt.title(title, size=30)
    plt.xlabel('身高 (cm)', size=15)
    plt.ylabel('体重 (kg)', size=15)
    plt.savefig(img.joinpath(title))


class KDTree:
    """
    高效算法KD树
    详细算法介绍：https://www.joinquant.com/view/community/detail/c2c41c79657cebf8cd871b44ce4f5d97
    """

    def __init__(self, point_list: List[Tuple[Tuple[int, int], str]],
                 depth: int = 0, root=None):
        """
        拿左右分支中的一半数据来与坐标点进行计算？
        :param point_list: Tuple[Tuple[x, y], 标签] 特征点
        :param depth: 深度
        :param root: 父节点
        """
        if len(point_list) > 0:
            # 轮换按照树深度选择坐标轴
            # 这里的k是特征数，也就是升高，体重
            k = len(point_list[0][0])
            # TODO 这个坐标是什么意思呢？用来干啥的？
            # axis 表示当前节点用于划分空间的坐标轴（维度）
            # 通过交替使用不同坐标轴（depth % k实现轮换），
            # 可以更均匀地划分多维空间，避免单一维度划分导致的搜索效率下降
            axis = depth % k

            # 选中位线，切 排序的依据为（身高，体重）
            point_list.sort(key=lambda x: x[0][axis])
            # 或者中间的点
            median = len(point_list) // 2
            # 当前划分轴
            self.axis = axis
            # 当前节点的父节点
            self.root = root

            # 造节点
            self.node = point_list[median]
            # 递归造左枝和右枝
            if len(point_list[:median]) > 0:
                self.left = KDTree(point_list[:median], depth + 1, self)
            else:
                self.left = None
            if len(point_list[median + 1:]) > 0:
                self.right = KDTree(point_list[median + 1:], depth + 1, self)
            else:
                self.right = None

    def find_leaf(self, point: Tuple[float, float]):
        """
        按切分寻找叶子 向左还是向右进查找，
        根据已知点的坐标轴与节点坐标轴进行比对大小 决定查找方向
        :param point: 坐标点
        :return:
        """
        # 最末尾的节点
        if self.left == None and self.right == None:
            return self
        # 左边无节点，向右边查找
        elif self.left == None:
            return self.right.find_leaf(point)
        # 右边无节点，向左边查找
        elif self.right == None:
            return self.left.find_leaf(point)
        # 根据当前节点的切分策略，来比对是向左查找还是向右查找
        elif point[self.axis] < self.node[0][self.axis]:
            return self.left.find_leaf(point)
        else:
            return self.right.find_leaf(point)

    def point_nearest(self, point: Tuple[float, float], k: int = 1, dist: Callable = L2_dist):
        """
        从最顶层开时，一层一层往下寻找最底层的点
        再从底层点往上找
        查找最近的k个点，复杂度 O(DlogN)，D是维度，N是树的大小
        :param point: 坐标点
        :param k: 距离
        :param dist: 距离函数，默认L2
        :return:
        """
        # 往下戳到最底叶
        leaf = self.find_leaf(point)
        # 从叶子网上爬 从最末支节点向上寻找
        return leaf.k_down_up(point, k, dist, result=list(), parent_node=self, child_node=None)

    def k_down_up(self, point: Tuple[float, float], k: int, dist: Callable,
                  result: List[Tuple[Tuple[Tuple[int, int], str], float]],
                  parent_node=None, child_node=None):
        """
        从下往上爬函数
        从最末支节点向上寻找
        当节点达到末支节点后，将其记录下来，如果记录表中不为空，
        :param point:
        :param k:
        :param dist:
        :param result:
        :param parent_node: 当前对象，节点到哪里去
        :param child_node: 当前对象，从哪里来
        :return: List[Tuple[Tuple[Tuple[int, int], str]]
        """
        # 选最长距离
        if len(result) == 0:
            max_dist = 0
        else:
            max_dist = max([x[1] for x in result])

        other_result = []

        # 从根节点开始找：
        #   目标节点与根节点的x坐标差值小于当前记录值的最大值 或 记录表中的点小于最大数量
        #   进入相应方向进行查找
        if (self.right == child_node and self.left is not None and
                ((point[self.axis] - self.node[0][self.axis] < max_dist) or
                 len(result) < k)):
            other_result = self.left.point_nearest(point, k, dist)
        if (self.left == child_node and self.right is not None and
                ((self.node[0][self.axis] - point[self.axis] < max_dist) or
                 len(result) < k)):
            other_result = self.right.point_nearest(point, k, dist)

        # 记录目标节点与当前根节点的距离
        result.append((self.node, dist(point, self.node[0])))
        # 取前k个结果
        result = sorted(result + other_result, key=lambda pair: pair[1])[:k]

        # 当前节点是否为根节点
        if self == parent_node:
            return result
        # 处于当前子分支中
        return self.root.k_down_up(point, k, dist, result, parent_node, self)

    def knn_prob(self, point: Tuple[float, float], label: str, k: int,
                 dist: Callable = L2_dist):
        """
        返回这个点属于该类别的概率
        :param point:
        :param label: 节点标签
        :param k:
        :param dist:
        :return:
        """
        nearests = self.point_nearest(point, k, dist)
        return float(len([pair for pair in nearests if pair[0][1] == label])) / float(len(nearests))

    def knn(self, point: Tuple[float, float], k: int, dist: Callable = L2_dist) -> Tuple[str, float]:
        """
        返回该点概率最大的类别以及相对应的概率
        :param point: 全局分布坐标点
        :param k: 最近距离
        :param dist: 归一化函数
        :return:
        """
        nearest = self.point_nearest(point, k, dist)

        # 分类出每个分类各有多少个
        statistics = {}
        for data in nearest:
            # data:Tuple[Tuple[x, y], 标签]
            label: str = data[0][1]
            if label not in statistics:
                statistics[label] = 1
            else:
                statistics[label] += 1
        # 找出数量最多的标签的标签
        max_label = max(statistics.items(), key=itemgetter(1))[0]
        return max_label, float(statistics[max_label]) / float(len(nearest))
